#include <algorithm/analytical_model/algorithm/cuda/Utils.h>
#include <algorithm/analytical_model/algorithm/cuda/evaluation/metrics/WRAD.cuh>

#include <cuda_runtime.h>
#include <thrust/async/scan.h>
#include <thrust/async/for_each.h>
#include <thrust/async/reduce.h>
#include <thrust/async/transform.h>
#include <thrust/device_ptr.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
#include <vector_functions.h>

#include <cmath>

using namespace solar::cuda;

auto solar::cuda::metrics::WRAD(float* d_input, float* d_ground_truth,
                                const inter_op::Int2& resolution,
                                const cudaStream_t stream) -> float
{
    auto result = metrics::async::WRAD(d_input, d_ground_truth, resolution, stream);
    return std::get<0>(result).get() / std::get<1>(result).get();
}

auto solar::cuda::metrics::async::WRAD(float* d_input, float* d_ground_truth,
                                       const inter_op::Int2& resolution, const cudaStream_t stream)
    -> std::tuple<thrust::cuda::unique_eager_future<float>,
                  thrust::cuda::unique_eager_future<float>>
{
    const std::size_t total_pixel_num = resolution.x * resolution.y;
    float* tmp_inclusive_scan_array = nullptr;
    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(tmp_inclusive_scan_array),
                                      total_pixel_num * sizeof(float), stream));
    const thrust::device_ptr<float> ground_truth_ptr(d_ground_truth);
    auto zip_iter_input_gt = thrust::make_zip_iterator(d_input, d_ground_truth);
    thrust::async::inclusive_scan(thrust::device.on(stream), d_ground_truth,
                                  d_ground_truth + total_pixel_num, tmp_inclusive_scan_array,
                                  thrust::plus<float>());

    float* result_diff = nullptr;
    float* per_pixel_weight = nullptr;
    float* weighted_ground_truth = nullptr;
    CHECK_CUDA_ERRORS(
        cudaMallocAsync(std::addressof(result_diff), total_pixel_num * sizeof(float), stream));
    CHECK_CUDA_ERRORS(
        cudaMallocAsync(std::addressof(per_pixel_weight), total_pixel_num * sizeof(float), stream));
    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(weighted_ground_truth),
                                      total_pixel_num * sizeof(float), stream));
    thrust::async::transform(thrust::device.on(stream), zip_iter_input_gt,
                             zip_iter_input_gt + total_pixel_num, result_diff,
                             [] __host__ __device__(const thrust::tuple<float, float>& iter)
                             { return std::abs(thrust::get<0>(iter) - thrust::get<1>(iter)); });
    thrust::async::transform(
        thrust::device.on(stream), d_ground_truth, d_ground_truth + total_pixel_num,
        per_pixel_weight,
        [tmp_inclusive_scan_array, total_pixel_num] __host__ __device__(const float& num)
        {
            const auto ground_truth_sum_raw = tmp_inclusive_scan_array[total_pixel_num - 1];
            return num / ground_truth_sum_raw;
        });
    auto zip_iter_gt_weight = thrust::make_zip_iterator(d_ground_truth, per_pixel_weight);
    auto zip_iter_diff_weight = thrust::make_zip_iterator(result_diff, per_pixel_weight);
    thrust::async::transform(thrust::device.on(stream), zip_iter_gt_weight,
                             zip_iter_gt_weight + total_pixel_num, weighted_ground_truth,
                             [] __host__ __device__(const thrust::tuple<float, float>& iter)
                             { return thrust::get<0>(iter) * thrust::get<1>(iter); });
    thrust::async::transform(thrust::device.on(stream), zip_iter_diff_weight,
                             zip_iter_diff_weight + total_pixel_num, result_diff,
                             [] __host__ __device__(const thrust::tuple<float, float>& iter)
                             { return thrust::get<0>(iter) * thrust::get<1>(iter); });
    auto weighted_diff_sum =
        thrust::async::reduce(thrust::device.on(stream), result_diff, result_diff + total_pixel_num,
                              0.0F, thrust::plus<float>());
    auto weighted_ground_truth_sum =
        thrust::async::reduce(thrust::device.on(stream), weighted_ground_truth,
                              weighted_ground_truth + total_pixel_num, 0.0F, thrust::plus<float>());
    CHECK_CUDA_ERRORS(cudaFreeAsync(result_diff, stream));
    CHECK_CUDA_ERRORS(cudaFreeAsync(per_pixel_weight, stream));
    CHECK_CUDA_ERRORS(cudaFreeAsync(weighted_ground_truth, stream));
    return {std::move(weighted_diff_sum), std::move(weighted_ground_truth_sum)};
}
